package com.example.harmonet.harmtorch;

public class Linear implements Layer {

    int col, row;
    float[][] weight;
    boolean hasBias;
    float[] bias;
    int param;

    public Linear(int input, int output) {
        this.col = input;
        this.row = output;
        this.weight = new float[row][col];
        this.bias = new float[row];
        this.hasBias = true;
        this.param = col * row + row;
    }

    public Linear(int input, int output, boolean hasBias) {
        this.col = input;
        this.row = output;
        this.weight = new float[row][col];
        this.hasBias = hasBias;
        if (this.hasBias) {
            this.bias = new float[row];
            this.param = col * row + row;
        }
        else {
            this.param = col * row;
        }
    }

    @Override
    public Tensor forward(Tensor in) throws Exception {
        if (in.dim()[in.dim().length - 1] != col) {
            throw new Exception("input dimension incompatible with linear");
        }
        int[] out_dim = in.dim();
        out_dim[out_dim.length - 1] = row;
        Tensor out = new Tensor(out_dim);
        int batch = 1;
        for(int i = 0; i < in.dim().length - 1; ++i) {
            batch *= in.dim()[i];
        }
        for(int n = 0; n < batch; ++n) {
            for(int i = 0; i < row; ++i) {
                out.tensor()[n * row + i] = bias[i];
                for(int j = 0; j < col; ++j) {
                    out.tensor()[n * row + i] += weight[i][j] * in.tensor()[n * col + j];
                }
            }
        }
        return out;
    }

    @Override
    public int getParam() {
        return param;
    }

    @Override
    public void init(float[] param) {
        for(int i = 0; i < row; ++i) {
            for(int j = 0; j < col; ++j) {
                this.weight[i][j] = param[i * col + j];
            }
        }
        for(int i = 0; i < row; ++i) {
            this.bias[i] = param[row * col + i];
        }
    }
}

